# -*- coding:UTF-8 -*-
import RPi.GPIO as GPIO
import time

# 小车电机引脚定义
IN1 = 20
IN2 = 21
IN3 = 19
IN4 = 26
ENA = 16
ENB = 13

# 小车按键定义
key = 8

# 设置GPIO口为BCM编码方式
GPIO.setmode(GPIO.BCM)

# 忽略警告信息
GPIO.setwarnings(False)


# 电机引脚初始化为输出模式
# 按键引脚初始化为输入模式
# 寻迹引脚初始化为输入模式
def init():
    global pwm_ENA
    global pwm_ENB
    GPIO.setup(ENA, GPIO.OUT, initial=GPIO.HIGH)
    GPIO.setup(IN1, GPIO.OUT, initial=GPIO.LOW)
    GPIO.setup(IN2, GPIO.OUT, initial=GPIO.LOW)
    GPIO.setup(ENB, GPIO.OUT, initial=GPIO.HIGH)
    GPIO.setup(IN3, GPIO.OUT, initial=GPIO.LOW)
    GPIO.setup(IN4, GPIO.OUT, initial=GPIO.LOW)
    GPIO.setup(key, GPIO.IN)
    # 设置pwm引脚和频率为2000hz
    pwm_ENA = GPIO.PWM(ENA, 2000)
    pwm_ENB = GPIO.PWM(ENB, 2000)
    pwm_ENA.start(0)
    pwm_ENB.start(0)


# 小车前进
def run(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.HIGH)
    GPIO.output(IN2, GPIO.LOW)
    GPIO.output(IN3, GPIO.HIGH)
    GPIO.output(IN4, GPIO.LOW)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车后退
def back(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.LOW)
    GPIO.output(IN2, GPIO.HIGH)
    GPIO.output(IN3, GPIO.LOW)
    GPIO.output(IN4, GPIO.HIGH)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车左转
def go_left(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.LOW)
    GPIO.output(IN2, GPIO.LOW)
    GPIO.output(IN3, GPIO.HIGH)
    GPIO.output(IN4, GPIO.LOW)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车右转
def go_right(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.HIGH)
    GPIO.output(IN2, GPIO.LOW)
    GPIO.output(IN3, GPIO.LOW)
    GPIO.output(IN4, GPIO.LOW)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车原地左转
def spin_left(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.LOW)
    GPIO.output(IN2, GPIO.HIGH)
    GPIO.output(IN3, GPIO.HIGH)
    GPIO.output(IN4, GPIO.LOW)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车原地右转
def spin_right(leftspeed, rightspeed):
    GPIO.output(IN1, GPIO.HIGH)
    GPIO.output(IN2, GPIO.LOW)
    GPIO.output(IN3, GPIO.LOW)
    GPIO.output(IN4, GPIO.HIGH)
    pwm_ENA.ChangeDutyCycle(leftspeed)
    pwm_ENB.ChangeDutyCycle(rightspeed)


# 小车停止
def brake():
    GPIO.output(IN1, GPIO.LOW)
    GPIO.output(IN2, GPIO.LOW)
    GPIO.output(IN3, GPIO.LOW)
    GPIO.output(IN4, GPIO.LOW)


import math
import cv2
import numpy as np


def trace_fun(imageFrame):
    """寻迹函数: 保持"""
    height = 100
    width = len(imageFrame[0])
    widthCenter = math.floor(width / 2)  # 320
    # 单单看图片的中间高度上一行的像素
    color = dst[math.floor(height)]  # 400
    # 找到白色的像素点个数
    white_count = np.sum(color == 255)
    white_index = np.where(color == 0)[0]
    arr = []
    if len(white_index > 0):
        arr.append(white_index[0])
        for i in range(len(white_index) - 1):
            if white_index[i + 1] - white_index[i] > 1:
                arr.append(white_index[i])
                arr.append(white_index[i + 1])
        arr.append(white_index[-1])
    return arr, height


cap = cv2.VideoCapture(1)


def getLines():
    left = None
    right = None
    for i in range(0, len(arr), 2):
        if arr[i + 1] - arr[i] <= blackLineWidth:
            if arr[i + 1] < mediumWeight:
                cv2.line(frame, (arr[i], height), (arr[i + 1], height), (0, 255, 0), 6, 4)
                left = (arr[i + 1] + arr[i]) / 2
            else:
                cv2.line(frame, (arr[i], height), (arr[i + 1], height), (255, 255, 255), 6, 4)
                right = (arr[i + 1] + arr[i]) / 2
    return left, right


while True:
    ret, frame = cap.imread()
    # 转化为灰度图
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # gray = cv2.equalizeHist(gray)
    dst = cv2.blur(gray, (2, 2))
    dst = cv2.medianBlur(dst, 3)
    # 大津法二值化
    retval, dst = cv2.threshold(dst, 0, 255, cv2.THRESH_OTSU)
    dst = cv2.medianBlur(dst, 9)
    dst = cv2.erode(dst, None, iterations=2)
    dst = cv2.blur(dst, (2, 2))

    arr, height = trace_fun(dst)
    weight = len(frame[0])
    blackLineWidth = math.floor(weight * 0.1)
    mediumWeight = math.floor(weight / 2)
    left, right = getLines()

    if left is not None and right is not None:
        run(35, 35)
        pass
    elif left is None and right is not None:
        spin_left(1, 1)
        pass
    elif left is not None and right is None:
        spin_right(1, 1)
        pass
    else:
        spin_right(1, 1)
        pass
    cv2.imshow("real", frame)
    cv2.imshow("dst", dst)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        cap.release()
        cv2.destroyAllWindows()
        break
